import mindspore.nn as nn


class LetNet5(nn.Cell):
    def __init__(self, num_classes=10):
        super(LetNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5, pad_mode='valid')
        self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Dense(16 * 5 * 5, 120)
        self.fc2 = nn.Dense(120, 84)
        self.fc3 = nn.Dense(84, num_classes)

        pass

    def construct(self, x):
        super(LetNet5, self).construct()
        x = self.conv1(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x


net = LetNet5()

epochs = 5
batch_size = 32
learning_rate = 0.01
momentum = .9

net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
net_opt = nn.Momentum(net.trainable_params(), learning_rate, momentum)

from mindspore import Model

model = Model(network=net, loss_fn=net_loss, optimizer=net_opt, metrics={
    "Accuracy": nn.Accuracy()
})

import t10

print(t10.aa)
# model.train(epochs, t10.dataset_train)

from mindvision.engine.callback import LossMonitor

model.train(epochs, t10.dataset_train, callbacks=[LossMonitor(.01, 1875)])

acc = model.eval(t10.dataset_eval, None, True)
print('{}', format(acc))

import mindspore as ms
import os


def mkdir(path):
    foler = os.path.exists(path)
    if not foler:
        os.makedirs(path)
        print("folder creat successfully", path)
    else:
        print('foler already exists', path)


file = r'./MyNet'
mkdir(file)

ms.save_checkpoint(net, './MyNet/save_direct.ckpt')
print('model has been saved')

# second method to keep the model
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig

config_ck = CheckpointConfig(save_checkpoint_steps=1875, keep_checkpoint_max=10)

ckpoint = ModelCheckpoint(prefix='lenet', directory='./MyNet/step_strategy', config=config_ck)
model.train(2, t10.dataset_train, callbacks=[ckpoint], dataset_sink_mode=False)

from mindvision.engine.callback import ValAccMonitor

model.train(3, t10.dataset_train, callbacks=
[ValAccMonitor(model, t10.dataset_eval, num_epochs=3
               , ckpt_directory='./MyNet')])

from mindspore import load_checkpoint, load_param_into_net
from mindvision.classification.dataset import Mnist
from mindvision.classification.models import lenet

param_dict = load_checkpoint('./MyNet/best.ckpt')
net = LetNet5()
load_param_into_net(net, param_dict)
net_opt = nn.Momentum(net.trainable_params(), learning_rate=.01, momentum=.9)
model = Model(network=net, loss_fn=net_loss, optimizer=net_opt,
              metrics={"Accuracy": nn.Accuracy()})

import numpy as np
from mindspore import Tensor
import matplotlib.pyplot as plt

mnist = Mnist('./datasets/mnist', split='test', batch_size=8, resize=32)
dataset_infer = mnist.run()
data_test = dataset_infer.create_dict_iterator()
data = next(data_test)

images = data['image'].asnumpy()
labels = data['label'].asnumpy()
plt.figure()
for i in range(1, 9):
    plt.subplot(2, 4, i)
    plt.imshow(images[i - 1].squeeze(), cmap='gray')
    plt.title(labels[i - 1])
plt.show()

output = model.predict(Tensor(data['image']))
predicted = np.argmax(output.asnumpy(), axis=1)

print(predicted, labels)
